#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
文本处理工具
包含文本到音素转换等功能
"""

import logging
import os
import sys
import re
import json
import time
from pathlib import Path

try:
    import nltk
    from g2p_en import G2p
    G2p_AVAILABLE = True
except ImportError:
    nltk = None
    G2p = None
    G2p_AVAILABLE = False

# 导入项目模块
sys.path.append(str(Path(__file__).resolve().parent.parent))
from config.paths import SYMBOL2ID_FILE, LEXICON_FILE, TEMP_DIR, NLTK_DATA_DIR, create_directories

# 配置日志
logger = logging.getLogger("文本处理")

class TextProcessor:
    """文本处理器类"""
    
    def __init__(self):
        """初始化文本处理器"""
        # 确保必要的目录结构存在
        create_directories()
        
        self.g2p = None
        
        # 设置NLTK数据目录
        if nltk:
            os.environ['NLTK_DATA'] = str(NLTK_DATA_DIR)
            nltk.data.path.insert(0, str(NLTK_DATA_DIR))
        
        # 自动下载NLTK数据
        if G2p_AVAILABLE:
            self._ensure_nltk_data()
            try:
                self.g2p = G2p()
                logger.info("G2P转换器已初始化")
            except Exception as e:
                logger.warning(f"G2P转换器初始化失败: {str(e)}")
                self._handle_nltk_error()
        else:
            logger.warning("G2P模块未安装，文本到音素转换功能将不可用")
            logger.warning("请安装g2p_en模块: pip install g2p_en")
        
        # 加载符号映射
        self.symbol2id = {}
        if os.path.exists(SYMBOL2ID_FILE):
            with open(SYMBOL2ID_FILE, 'r') as f:
                self.symbol2id = json.load(f)
        
        # 加载发音词典
        self.lexicon = {}
        if os.path.exists(LEXICON_FILE):
            with open(LEXICON_FILE, 'r') as f:
                for line in f:
                    parts = line.strip().split()
                    if len(parts) > 1:
                        word = parts[0]
                        phones = ' '.join(parts[1:])
                        self.lexicon[word] = phones
    
    def _ensure_nltk_data(self):
        """确保NLTK数据已下载"""
        if not nltk:
            logger.warning("NLTK未安装，无法使用文本到音素转换功能")
            return False
        
        try:
            # 设置NLTK只使用项目内的数据目录
            os.environ['NLTK_DATA'] = str(NLTK_DATA_DIR)
            nltk.data.path = [str(NLTK_DATA_DIR)]
            
            # 直接检查文件是否存在，而不是使用nltk.data.find()
            required_files = {
                'cmudict': NLTK_DATA_DIR / "corpora" / "cmudict" / "cmudict",
                'averaged_perceptron_tagger': NLTK_DATA_DIR / "taggers" / "averaged_perceptron_tagger" / "averaged_perceptron_tagger.pickle"
            }
            
            all_available = True
            missing_files = []
            
            # 检查必要文件
            for name, file_path in required_files.items():
                if not file_path.exists():
                    all_available = False
                    missing_files.append(name)
                    logger.warning(f"缺少NLTK资源: {name}")
            
            # 如果文件缺失，尝试下载
            if not all_available:
                logger.warning(f"缺少NLTK数据文件: {', '.join(missing_files)}")
                self._handle_nltk_error()
                
                # 再次检查文件是否存在
                all_available = True
                for name, file_path in required_files.items():
                    if not file_path.exists():
                        all_available = False
                        logger.error(f"下载后仍缺少NLTK资源: {name}")
                        
            # 检查g2p_en所需的额外资源
            if all_available and G2p_AVAILABLE:
                try:
                    from main import prepare_g2p_resources
                    prepare_g2p_resources()
                except Exception as e:
                    logger.error(f"准备g2p资源时出错: {str(e)}")
            
            return all_available
        except Exception as e:
            logger.error(f"检查NLTK数据时出错: {str(e)}")
            return False
    
    def _handle_nltk_error(self):
        """处理NLTK数据下载错误"""
        error_msg = (
            "NLTK数据下载失败。您可以尝试以下方法解决：\n"
            "1. 从「工具」菜单使用「NLTK数据下载」功能\n"
            "2. 或者手动运行: python -m nltk.downloader -d ./data/nltk_data cmudict averaged_perceptron_tagger\n"
            "3. 检查网络连接是否正常\n"
        )
        logger.warning(error_msg)
        
        # 尝试自动下载并保存NLTK数据文件
        self._try_manual_nltk_download()
    
    def _try_manual_nltk_download(self):
        """尝试手动方式下载NLTK数据"""
        try:
            # 创建NLTK数据目录
            os.makedirs(NLTK_DATA_DIR / "corpora", exist_ok=True)
            os.makedirs(NLTK_DATA_DIR / "taggers", exist_ok=True)
            
            # 下载数据（如果有互联网连接）
            import urllib.request
            
            # 下载超时时间（秒）
            timeout = 10
            
            # 下载CMU词典
            logger.info("尝试手动下载CMU词典...")
            cmudict_url = "https://raw.githubusercontent.com/nltk/nltk_data/gh-pages/packages/corpora/cmudict.zip"
            cmudict_file = f"{TEMP_DIR}/cmudict.zip"
            urllib.request.urlretrieve(cmudict_url, cmudict_file, timeout=timeout)
            
            # 下载词性标注器
            logger.info("尝试手动下载词性标注器...")
            tagger_url = "https://raw.githubusercontent.com/nltk/nltk_data/gh-pages/packages/taggers/averaged_perceptron_tagger.zip"
            tagger_file = f"{TEMP_DIR}/averaged_perceptron_tagger.zip"
            urllib.request.urlretrieve(tagger_url, tagger_file, timeout=timeout)
            
            # 解压数据
            import zipfile
            
            # 解压CMU词典
            logger.info("正在解压CMU词典...")
            with zipfile.ZipFile(cmudict_file, 'r') as zip_ref:
                zip_ref.extractall(NLTK_DATA_DIR / "corpora")
            
            # 解压词性标注器
            logger.info("正在解压词性标注器...")
            with zipfile.ZipFile(tagger_file, 'r') as zip_ref:
                zip_ref.extractall(NLTK_DATA_DIR / "taggers")
            
            logger.info("NLTK数据手动下载并解压完成")
            
            # 小延迟，确保文件已解压完成
            time.sleep(1)
            
            # 删除临时文件
            os.remove(cmudict_file)
            os.remove(tagger_file)
            
            return True
        except Exception as e:
            logger.warning(f"手动下载NLTK数据失败: {str(e)}")
            return False
    
    def text_to_phonemes(self, text):
        """
        将文本转换为音素序列
        
        Args:
            text: 要转换的文本
            
        Returns:
            音素序列字符串
        """
        # 预处理文本
        text = self._preprocess_text(text)
        
        # 如果G2P未初始化，尝试初始化
        if self.g2p is None and G2p_AVAILABLE:
            try:
                self._ensure_nltk_data()
                self.g2p = G2p()
                logger.info("G2P转换器已初始化")
            except Exception as e:
                logger.warning(f"G2P转换器初始化失败: {str(e)}")
                # g2p将保持为None
        
        # 尝试使用g2p转换
        if self.g2p is not None:
            try:
                # 转换为音素
                phonemes = self.g2p(text)
                
                # 后处理音素序列
                phoneme_sequence = self._postprocess_phonemes(phonemes)
                
                return phoneme_sequence
            except Exception as e:
                logger.error(f"文本转换为音素时出错: {str(e)}")
                # 继续执行备用方法
        
        # 备用方法：简单的英文字母到音素映射
        logger.warning("使用备用的音素转换方法")
        simple_phonemes = []
        
        # 简单的映射表（确保所有值都有@前缀）
        simple_map = {
            'a': '@AE1', 'e': '@EH1', 'i': '@IH1', 'o': '@OW1', 'u': '@AH1',
            'b': '@B', 'c': '@K', 'd': '@D', 'f': '@F', 'g': '@G',
            'h': '@HH', 'j': '@JH', 'k': '@K', 'l': '@L', 'm': '@M',
            'n': '@N', 'p': '@P', 'q': '@K', 'r': '@R', 's': '@S',
            't': '@T', 'v': '@V', 'w': '@W', 'x': '@K @S', 'y': '@Y', 'z': '@Z'
        }
        
        # 转换每个字符
        for char in text.lower():
            if char in simple_map:
                if ' ' in simple_map[char]:  # 处理多音素情况，如'x'映射到'@K @S'
                    simple_phonemes.extend(simple_map[char].split())
                else:
                    simple_phonemes.append(simple_map[char])
            elif char.isalpha():
                # 未知字母用"@sp"（特殊音素）表示
                simple_phonemes.append('@sp')
            elif char.isspace():
                # 空格用静音表示
                simple_phonemes.append('@SILENCE')
        
        # 如果没有任何有效音素，返回占位符
        if not simple_phonemes:
            simple_phonemes = ['@sp']
            
        return ' '.join(simple_phonemes)
    
    def _preprocess_text(self, text):
        """
        预处理文本
        
        Args:
            text: 原始文本
            
        Returns:
            预处理后的文本
        """
        # 移除多余空格
        text = re.sub(r'\s+', ' ', text.strip())
        
        # 转换为小写
        text = text.lower()
        
        # 这里可以添加更多预处理步骤，如标点符号处理等
        
        return text
    
    def _postprocess_phonemes(self, phonemes):
        """
        处理音素序列
        
        Args:
            phonemes: 音素列表
            
        Returns:
            处理后的音素序列字符串
        """
        # 移除音素中的数字（音调标记）
        cleaned_phonemes = []
        for p in phonemes:
            if p == ' ':
                cleaned_phonemes.append(p)
            elif p == '{':
                cleaned_phonemes.append('.')
            else:
                # 保留音调标记(0,1,2)，因为模型需要区分重音
                # 添加@前缀以匹配symbol2id.json格式
                if p:
                    # 确保音素有@前缀
                    if not p.startswith('@'):
                        cleaned_phonemes.append(f"@{p}")
                    else:
                        cleaned_phonemes.append(p)
        
        # 拼接成字符串
        phoneme_sequence = ' '.join(cleaned_phonemes)
        
        return phoneme_sequence
    
    def phonemes_to_ids(self, phonemes):
        """
        将音素序列转换为ID序列
        
        Args:
            phonemes: 音素序列（字符串或列表）
            
        Returns:
            ID序列列表
        """
        if isinstance(phonemes, str):
            phonemes = phonemes.split()
        
        ids = []
        unknown_phonemes = set()  # 用于跟踪未知音素
        
        for p in phonemes:
            # 跳过空格
            if p.strip() == '':
                continue
            
            # 确保音素有@前缀
            if p and not p.startswith('@'):
                p = f"@{p}"
            
            # 查找音素ID
            if p in self.symbol2id:
                ids.append(self.symbol2id[p])
            else:
                if p not in unknown_phonemes:
                    unknown_phonemes.add(p)
                    logger.warning(f"未知的音素: {p}")
                
                # 使用特殊符号
                if '@sp' in self.symbol2id:
                    ids.append(self.symbol2id['@sp'])
                elif '@BLANK' in self.symbol2id:
                    ids.append(self.symbol2id['@BLANK'])
                else:
                    # 默认使用PAD符号
                    ids.append(0)
        
        # 确保至少有一个ID
        if not ids:
            if '@SILENCE' in self.symbol2id:
                ids.append(self.symbol2id['@SILENCE'])
            else:
                ids.append(0)  # 使用PAD符号
                
        # 记录转换结果
        logger.debug(f"音素转换为ID: {phonemes} -> {ids}")
            
        return ids
    
    def normalize_text(self, text):
        """
        文本归一化
        
        Args:
            text: 输入文本
            
        Returns:
            归一化后的文本
        """
        # 移除多余空格
        text = re.sub(r'\s+', ' ', text.strip())
        
        # 转换为小写
        # text = text.lower()
        
        # 更多文本归一化规则可以在这里添加
        
        return text

# 测试代码
if __name__ == "__main__":
    processor = TextProcessor()
    
    test_text = "Hello, this is a test sentence with 123 numbers!"
    print(f"原始文本: {test_text}")
    
    phonemes = processor.text_to_phonemes(test_text)
    print(f"音素序列: {phonemes}")
    
    ids = processor.phonemes_to_ids(phonemes)
    print(f"ID序列: {ids}") 